#!/usr/bin/python
# -*- coding: utf-8 -*-
from __future__ import division
import sys, argparse
import csv
import time as t
import numpy as np
import random
import datetime as d
import pandas as pd
'''
将训练集和测试集拆分成10份
'''

def parse_args():
    if len(sys.argv) == 1:
        sys.argv.append('-h')

    parser = argparse.ArgumentParser()
    parser.add_argument('enrollment')
    parser.add_argument('n_parts')
    parser.add_argument('output')
    args = vars(parser.parse_args())
    return args

args = parse_args()
enrollment, n_parts, output = args['enrollment'], args['n_parts'], args['output']

data = pd.read_csv(enrollment)

enrollment_ids = data.enrollment_id.values

random.shuffle(enrollment_ids)

n_parts = int(n_parts)


part_num_rcds = int(len(enrollment_ids) / n_parts)

for i in range(n_parts):
    output_path = "%s.%d" % (output, i)
    part = enrollment_ids[i * part_num_rcds: (i+1) * part_num_rcds]
    with open(output_path, 'w') as f:
        for p in part:
            f.write(str(p) + '\n')
